# ---
# jupyter:
#   jupytext:
#     formats: py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.14.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %%
# remove-cell
# pylint: disable=line-too-long,redefined-outer-name

# %% [markdown]
# # Classification, other data distributions, VGP and SVGP
#
# In this chapter we will talk about what you can do if your data is not normally distributed.

# %% [markdown]
# As usual we will start with our imports:

# %%
# hide: begin
import os
import warnings

warnings.simplefilter("ignore")
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
# hide: end

from typing import Sequence

import matplotlib.pyplot as plt
import numpy as np

import gpflow

# hide: begin
# %matplotlib inline
plt.rcParams["figure.figsize"] = (12, 6)
# hide: end

# %% [markdown]
# ## The Variational Gaussion Process
#
# Remember how we assume that our data is generated by:
#
# \begin{equation}
# Y_i = f(X_i) + \varepsilon_i \,, \quad \varepsilon_i \sim \mathcal{N}(0, \sigma^2).
# \end{equation}
#
# In this chapter we will no longer assume this &mdash; instead we will allow any stochastic function of $f$.
#
# We will be using the Variational Gaussian Process (VGP) model in this chapter. The VGP uses a [Likelihood](../../api/gpflow/likelihoods/index.rst#gpflow-likelihoods-likelihood) object to define the connection between $f$ and $Y$. The models we have seen so far, the GPR and SGPR, requires the likelihood to be gaussian, but the VGP allows us to configure the likelihood.
#
#

# %% [markdown]
# ### Non-gaussian regression
#
# Let us first define a helper function, for plotting a model:

# %%
def plot_model(model: gpflow.models.GPModel) -> None:
    X, Y = model.data
    opt = gpflow.optimizers.Scipy()
    opt.minimize(model.training_loss, model.trainable_variables)
    gpflow.utilities.print_summary(model, "notebook")

    Xplot = np.linspace(0.0, 1.0, 200)[:, None]

    y_mean, y_var = model.predict_y(Xplot, full_cov=False)
    y_lower = y_mean - 1.96 * np.sqrt(y_var)
    y_upper = y_mean + 1.96 * np.sqrt(y_var)

    _, ax = plt.subplots(nrows=1, ncols=1)
    ax.plot(X, Y, "kx", mew=2)
    (mean_line,) = ax.plot(Xplot, y_mean, "-")
    color = mean_line.get_color()
    ax.plot(Xplot, y_lower, lw=0.1, color=color)
    ax.plot(Xplot, y_upper, lw=0.1, color=color)
    ax.fill_between(
        Xplot[:, 0], y_lower[:, 0], y_upper[:, 0], color=color, alpha=0.1
    )


# %% [markdown]
# Below we have a small dataset with an outlier. The default Gaussian likelihood has very light tails, and will struggle with the outlier:

# %%
# remove-cell

rng = np.random.default_rng(1235)
n = 10
X = rng.random((n, 1))
f = np.exp(X)
e = 0.05 * rng.standard_normal(X.shape)
Y = f + e
print("----------------------------------------------")
for x in X:
    print(f"[{float(x): .3f}],")
print("----------------------------------------------")
for y in Y:
    print(f"[{float(y): .2f}],")

# %%
# hide: begin
# fmt: off
# hide: end
X = np.array(
    [
        [0.177], [0.183], [0.428], [0.838], [0.827], [0.293], [0.270], [0.593],
        [0.031], [0.650],
    ]
)
Y = np.array(
    [
        [1.22], [1.17], [1.99], [2.29], [2.29], [1.28], [1.20], [1.82], [1.01],
        [1.93],
    ]
)
# hide: begin
# fmt: on
# hide: end

_ = plt.scatter(X, Y)

# %% [markdown]
# Let us try fitting the usual GPR to this data:

# %%
model: gpflow.models.GPModel = gpflow.models.GPR(
    (X, Y), kernel=gpflow.kernels.SquaredExponential()
)
plot_model(model)

# %% [markdown]
# Notice the high variance predicted by the model, and how the model consistently overestimate $f$ around the non-outlier middle points.
#
# Instead, let us try fitting a VGP with a [Student's t-distribution](https://en.wikipedia.org/wiki/Student%27s_t-distribution) as likelihood. The Student's t-distribution has heavier tails, and will be less affected by an outlier.

# %%
model = gpflow.models.VGP(
    (X, Y),
    kernel=gpflow.kernels.SquaredExponential(),
    likelihood=gpflow.likelihoods.StudentT(),
)
plot_model(model)

# %% [markdown]
# Notice how this fit has smaller uncertainty, and is a much better fit for the non-outlier points.

# %% [markdown]
# ### Classification
#
# Variational models and custom likelihoods allow you to have arbitrary connections beween the data you are modelling and the real-valued function $f$. For example, we can use it for a classification problem where we are not trying to predict a real value (regression), but trying to predict an object label instead.
#
# To do classification we start with our unknown function $f$. $f$ can have any (real) value, but we want to use it for probabilities, so we need to map it to the interval $[0; 1]$. We will use the cumulative density function of a standard normal for this. Other functions can be used for this "squishing", but the cumulative density function has nice analytical properties that simplifies some maths. We then use the mapped value as the probability in a [Bernoulli trial](https://en.wikipedia.org/wiki/Bernoulli_trial) to determine the output class ($Y$). All of this is encapsulated in the GPflow [Bernoulli](../../api/gpflow/likelihoods/index.rst#gpflow-likelihoods-bernoulli) class.

# %% [markdown]
# Assume we have two sets of points:

# %%
# remove-cell

rng = np.random.default_rng(1235)
n1 = 10
n2 = 10
X1 = rng.normal(0.4, 0.2, n1)
X2 = rng.normal(0.8, 0.15, n2)

plt.scatter(X1, np.zeros_like(X1))
plt.scatter(X2, np.zeros_like(X2))

# %%
# remove-cell

print("----------------------------------------------")
for x in X1:
    print(f"[{float(x): .3f}],")

print("----------------------------------------------")
for x in X2:
    print(f"[{float(x): .3f}],")

# %%
# hide: begin
# fmt: off
# hide: end
X1 = np.array(
    [[0.218], [0.453], [0.196], [0.638], [0.523], [0.541], [0.455], [0.632], [0.309], [0.330]]
)
X2 = np.array(
    [[0.868], [0.706], [0.672], [0.742], [0.813], [0.617], [0.456], [0.838], [0.730], [0.841]]
)
# hide: begin
# fmt: on
# hide: end

plt.scatter(X1, np.zeros_like(X1))
_ = plt.scatter(X2, np.zeros_like(X2))

# %% [markdown]
# Given a new point, we want to predict which one of those two sets it should belong to.
#
# First we will massage our data into one $(X, Y)$ dataset, instead of two $X$ datasets.
# Our $X$ will be all our points, and our $Y$ will be either 0 or 1, to signal which one of the two sets the corresponding $x$ belongs to.

# %%
Y1 = np.zeros_like(X1)
Y2 = np.ones_like(X2)
X = np.concatenate([X1, X2], axis=0)
Y = np.concatenate([Y1, Y2], axis=0)

_ = plt.scatter(X, Y)

# %% [markdown]
# Notice that these are the same points, but instead of using colours, we are using the Y-axis to separate the two sets.

# %% [markdown]
# We can train a model on this the usual way:

# %%
model = gpflow.models.VGP(
    (X, Y),
    kernel=gpflow.kernels.SquaredExponential(),
    likelihood=gpflow.likelihoods.Bernoulli(),
)
opt = gpflow.optimizers.Scipy()
opt.minimize(model.training_loss, model.trainable_variables)
gpflow.utilities.print_summary(model, "notebook")

# %% [markdown]
# When doing predictions, beware that since we are no longer using a Gaussian likelihood - in fact our likelihood is not even symmetric - it is much harder to make sense of the mean and variance returned by `predict_f` and `predict_y`. Instead, let us draw samples of f:

# %%
Xplot = np.linspace(0, 1, 200)[:, None]
Fsamples = model.predict_f_samples(Xplot, 10).numpy().squeeze().T

_ = plt.plot(Xplot, Fsamples, "C0", lw=0.5)

# %% [markdown]
# How do we interpret this? Remember these are $f$ samples from before they're "squished" into the $[0; 1]$ range. The `likelihood.invlink` function implements the above-mentioned cumulative density function that is used to perform the "squishing". We can use this to map our $f$ samples to probabilities:

# %%
Psamples = model.likelihood.invlink(Fsamples)

plt.plot(Xplot, Psamples, "C1", lw=0.5)
_ = plt.scatter(X, Y)

# %% [markdown]
# We can use the mean, "squished" through our `invlink` to get a probability of a point being in group 0:

# %%
Fmean, _ = model.predict_f(Xplot)
P = model.likelihood.invlink(Fmean)

plt.plot(Xplot, P, "C1")
_ = plt.scatter(X, Y)

# %% [markdown]
# So if we want to predict the class at $x=0.3$ we would do:

# %%
Xnew = np.array([[0.3]])
Fmean, _ = model.predict_f(Xnew)
P = model.likelihood.invlink(Fmean)
P

# %% [markdown]
# Which mean a $\sim3.7\%$ chance of class 0, and $100\% - 3.7\% = 96.3\%$ chance of class 1.

# %% [markdown]
# You may also be interested in our advanced [tutorial on multiclass classification](../advanced/multiclass_classification.ipynb).

# %% [markdown]
# ## The Sparse Variational Gaussian Process
#
# The Sparse Variational Gaussian Process (SVGP) combines the sparsity we studied in the previous chapter, with the generic likelihoods we have seen in this chapter. It gets a separate section, because its API is slightly different from what we've seen before. All the other models we have seen use the [InternalDataTrainingLossMixin](../../api/gpflow/models/index.rst#gpflow.models.InternalDataTrainingLossMixin) and internally store their data. The SVGP uses the [ExternalDataTrainingLossMixin](../../api/gpflow/models/index.rst#gpflow-models-externaldatatraininglossmixin). This means it does not take its data in the constructor, and instead takes the data as parameters when training the model.
#
# Let us try fitting a SVGP to our classification data:

# %%
model = gpflow.models.SVGP(
    kernel=gpflow.kernels.SquaredExponential(),
    likelihood=gpflow.likelihoods.Bernoulli(),
    inducing_variable=np.linspace(0.0, 1.0, 4)[:, None],
)
opt = gpflow.optimizers.Scipy()
opt.minimize(model.training_loss_closure((X, Y)), model.trainable_variables)
gpflow.utilities.print_summary(model, "notebook")

# %% [markdown]
# Notice the parameters provided to the model initialiser, and the parameters passed to `model.training_loss_closure`.
#
# Let us plot our results:

# %%
Fsamples = model.predict_f_samples(Xplot, 10).numpy().squeeze().T
Psamples = model.likelihood.invlink(Fsamples)

plt.plot(Xplot, Psamples, "C1", lw=0.5)
_ = plt.scatter(X, Y)


# %% [markdown]
# The combination of external training data and inducing points allow the SVGP to be trained on truly huge datasets. For details, see our [advanced tutorial on scalability](../advanced/gps_for_big_data.ipynb).

# %% [markdown]
# ### Writing code that handles both internal and external data models
#
# Having some models store training data internally, while others take it as a parameter when training, can be frustrating if you are trying to write code that works with a generic model. The `gpflow.models` module contains [some functions that can help you write generic code](../../api/gpflow/models/index.rst#functions).
#
# Here is an example of how one might write a generic model training function:

# %%
def train_generic_model(
    model: gpflow.models.GPModel, data: gpflow.base.RegressionData
) -> None:
    loss = gpflow.models.training_loss_closure(model, data)
    opt = gpflow.optimizers.Scipy()
    opt.minimize(loss, model.trainable_variables)


# %% [markdown]
# For example, this works with both a VGP and an SVGP:

# %%
data = (X, Y)
models: Sequence[gpflow.models.GPModel] = [
    gpflow.models.VGP(
        data,
        kernel=gpflow.kernels.SquaredExponential(),
        likelihood=gpflow.likelihoods.Bernoulli(),
    ),
    gpflow.models.SVGP(
        kernel=gpflow.kernels.SquaredExponential(),
        likelihood=gpflow.likelihoods.Bernoulli(),
        inducing_variable=np.linspace(0.0, 1.0, 4)[:, None],
    ),
]
for model in models:
    train_generic_model(model, data)
